package rabbit.flt.plugins.common;

import rabbit.flt.common.trace.TraceData;
import rabbit.flt.plugins.common.plugin.SupportPlugin;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.lang.reflect.Method;

import static rabbit.flt.common.trace.MessageType.R2DBC;

public abstract class R2dbcSupportPlugin extends SupportPlugin {

    /**
     * 手动创建traceData
     *
     * @param method
     * @param sql
     * @return
     */
    protected TraceData manuallyCreateTraceData(Method method, String sql) {
        TraceData traceData = createTraceData(method);
        traceData.setData(sql);
        traceData.setNodeName(getNameBySql(sql));
        traceData.setMessageType(R2DBC.name());
        return traceData;
    }

    /**
     * 包装结果
     *
     * @param result
     * @param method
     * @param sql
     * @return
     */
    protected Mono wrapResult(Mono<?> result, Method method, String sql) {
        return Mono.defer(() -> {
            TraceData traceData = manuallyCreateTraceData(method, sql);
            traceData.setRequestTime(System.currentTimeMillis());
            return result.doOnError(e -> setErrorInfo(traceData, e)).doOnSuccess(s -> traceData.updateCost()).map(s -> {
                traceData.updateCost();
                return s;
            }).doFinally(f -> super.handleTraceData(traceData));
        });
    }

    /**
     * 包装结果
     *
     * @param result
     * @return
     */
    protected Flux wrapResult(Flux<?> result, Method method, String sql) {
        return Flux.defer(() -> {
            TraceData traceData = manuallyCreateTraceData(method, sql);
            traceData.setRequestTime(System.currentTimeMillis());
            return result.doOnError(e -> setErrorInfo(traceData, e)).doOnComplete(traceData::updateCost).map(s -> {
                traceData.updateCost();
                return s;
            }).doFinally(f -> super.handleTraceData(traceData));
        });
    }

    private void setErrorInfo(TraceData traceData, Throwable e) {
        traceData.updateCost();
        traceData.setStatus(TraceData.Status.ERR);
        traceData.setData(getStackText(e));
    }
}
